import os
import csv
import json

# Grid size 100*100: 963160
def node_num():
    maxval = 0
    for i in range(400):
        filename = './grid/' + "ID-" + str(i) + '.csv'
        if not os.path.exists(filename):
            pass
        else:
            with open(filename, "r") as f:
                reader = csv.reader(f)
                for trajectories in reader:
                    for node in trajectories:
                        tmp = int(node)
                        if tmp > maxval:
                            maxval = tmp
    
    return maxval


def node_count_train2():
    nodes = [0]*963161
    filename = './ngram_statistics/size_100/2MMC/train4000k.csv'
    filename2 = './ngram_statistics/size_100/2MMC/node_count_train4000k.csv'
        
    with open(filename, "r") as f:
        reader = csv.reader(f)
        for trajectories in reader:
            nodes[int(trajectories[0])] += 1

    with open(filename2, 'w', newline='') as csvfile:
        spamwriter = csv.writer(csvfile)
        index = -1
        for item in nodes:
            index += 1
            if item > 0:
                spamwriter.writerow([index, item])


def node_count_train3():
    trigram = {}
    filename = './ngram_statistics/size_100/3MMC/train1000k.csv'
    filename2 = './ngram_statistics/size_100/3MMC/node_count_train1000k.json'
        
    with open(filename, "r") as f:
        reader = csv.reader(f)
        for trajectories in reader:
            prev = trajectories[0]
            cur = trajectories[1]
            k = prev+'_'+cur
            if not trigram.__contains__(k):
                trigram[k] = {'CNT':0}
            trigram[k]['CNT'] += 1

    json_data = json.dumps(trigram)
    with open(filename2, 'w') as f3:
        f3.write(json_data)


def node_pair3(n):
    with open('./ngram_statistics/size_100/3MMC/node_pair.csv', 'w', newline='') as csvfile:
        spamwriter = csv.writer(csvfile)
        for i in range(400):
            filename = './grid/ID-' + str(i) + '.csv'
            if not os.path.exists(filename):
                pass
            else:
                with open(filename, "r") as f:
                    reader = csv.reader(f)
                    for trajectories in reader:
                        if len(trajectories) < n:
                            pass
                        for i in range(len(trajectories)-2):
                            spamwriter.writerow([int(trajectories[i]), int(trajectories[i+1]), int(trajectories[i+2])])
                        spamwriter.writerow([int(trajectories[len(trajectories)-2]), int(trajectories[len(trajectories)-1]), 0])

def node_pair2(n):
    with open('./ngram_statistics/size_100/node_pair2.csv', 'w', newline='') as csvfile:
        spamwriter = csv.writer(csvfile)
        for i in range(400):
            filename = './grid/' + "ID-" + str(i) + '.csv'
            if not os.path.exists(filename):
                pass
            else:
                with open(filename, "r") as f:
                    reader = csv.reader(f)
                    for trajectories in reader:
                        if len(trajectories) < n:
                            pass
                        for i in range(len(trajectories)-1):
                            spamwriter.writerow([int(trajectories[i]), int(trajectories[i+1])])
                        spamwriter.writerow([int(trajectories[len(trajectories)-1]), 0])

def bigram():
    prob_matrix = {}
    node_index = []
    with open("./ngram_statistics/size_100/2MMC/node_count_train4000k.csv", "r") as f1:
        reader = csv.reader(f1)
        for item in reader:
            # nodeID = int(item[0])
            nodeCNT = int(item[1])
            prob_matrix[item[0]] = {}
            prob_matrix[item[0]]['CNT'] = nodeCNT
    
    with open("./ngram_statistics/size_100/2MMC/train4000k.csv", "r") as f2:
        reader = csv.reader(f2)
        for row in reader:
            if prob_matrix[row[0]].__contains__(row[1]):
                prob_matrix[row[0]][row[1]] += 1
            else:
                prob_matrix[row[0]][row[1]] = 1
    
    json_data = json.dumps(prob_matrix)
    
    with open('./ngram_statistics/size_100/2MMC/bigram_4000k.json', 'w') as f3:
        f3.write(json_data)

def trigram():
    with open('./ngram_statistics/size_100/3MMC/node_count_train1000k.json') as bf:
        prob_matrix = json.load(bf)

    node_index = []

    with open("./ngram_statistics/size_100/3MMC/train1000k.csv", "r") as f2:
        reader = csv.reader(f2)
        for row in reader:
            prev = row[0]
            cur = row[1]
            k = prev+'_'+cur
            if not prob_matrix.__contains__(k):
                continue
            if prob_matrix[k].__contains__(row[2]):
                prob_matrix[k][row[2]] += 1
            else:
                prob_matrix[k][row[2]] = 1
    
    json_data = json.dumps(prob_matrix)
    
    with open('./ngram_statistics/size_100/3MMC/trigram_1000k.json', 'w') as f3:
        f3.write(json_data)
            


if __name__ == '__main__':
    #node_pair3(3)
    node_count_train3()
    trigram()